#python paddleocr_basic.py Test.png
# 监控时间戳识别专用解决方案 - PaddleOCR实现
import os
import sys
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import cv2
from datetime import datetime
import json

def init_paddleocr():
    """初始化PaddleOCR，使用适合时间戳识别的配置"""
    try:
        from paddleocr import PaddleOCR
        
        # 设置环境变量避免版本冲突
        os.environ['FLAGS_allocator_strategy'] = 'naive_best_fit'
        os.environ['FLAGS_use_gpu'] = 'false'  # 强制使用CPU
        
        print("正在初始化PaddleOCR（时间戳识别专用配置）...")
        
        # 使用最稳定的配置
        ocr = PaddleOCR(
            use_angle_cls=True,  # 启用方向分类，适合监控图像
            lang='ch',  # 中文模式
            det_model_dir=None,  # 使用默认检测模型
            rec_model_dir=None,  # 使用默认识别模型
            cls_model_dir=None,  # 使用默认分类模型
            use_gpu=False,  # 强制CPU模式
            precision='fp32',  # 使用fp32精度，更稳定
            det_limit_side_len=960,  # 检测图像边长限制
            det_limit_type='max',  # 限制类型
            show_log=False  # 关闭详细日志
        )
        
        print("PaddleOCR初始化成功！")
        return ocr, True
        
    except ImportError as e:
        print(f"PaddleOCR未安装: {e}")
        print("请安装: pip install paddleocr")
        return None, False
    except Exception as e:
        print(f"PaddleOCR初始化失败: {e}")
        print("尝试使用基础配置...")
        
        try:
            # 尝试最简配置
            ocr = PaddleOCR(lang='ch', use_gpu=False)
            print("PaddleOCR基础配置初始化成功！")
            return ocr, True
        except Exception as e2:
            print(f"基础配置也失败: {e2}")
            return None, False

def preprocess_timestamp_image(image_path):
    """
    专门针对时间戳的图像预处理 - 使用直方图均衡化
    监控图像中时间戳通常特点：
    1. 位于图像角落
    2. 字体相对较小
    3. 对比度可能不高
    """
    try:
        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            return None, None
            
        original = image.copy()
        
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 直方图均衡化 - 改善对比度，让文字更清晰
        processed = cv2.equalizeHist(gray)
        
        return original, processed
        
    except Exception as e:
        print(f"图像预处理失败: {e}")
        return None, None

def filter_timestamp_results(results, min_confidence=0.2):
    """
    过滤和优化时间戳识别结果 - 增强版，不依赖固定格式
    """
    if not results or not results[0]:
        return []
    
    filtered_results = []
    
    for line in results[0]:
        if not line:
            continue
            
        bbox, (text, confidence) = line[0], line[1]
        
        # 置信度过滤（降低阈值，允许更多候选）
        if confidence < min_confidence:
            continue
            
        # 清理文本
        cleaned_text = text.strip()
        if not cleaned_text or len(cleaned_text) < 2:
            continue
            
        # 增强的时间戳相关性评分（不依赖固定格式）
        timestamp_score = calculate_enhanced_timestamp_score(cleaned_text)
        
        filtered_results.append({
            'text': cleaned_text,
            'confidence': confidence,
            'bbox': bbox,
            'timestamp_score': timestamp_score,
            'total_score': confidence * 0.6 + timestamp_score * 0.4,  # 调整权重
            'category': classify_timestamp_text(cleaned_text)
        })
    
    # 按综合得分排序
    filtered_results.sort(key=lambda x: x['total_score'], reverse=True)
    
    return filtered_results

def calculate_enhanced_timestamp_score(text):
    """计算增强的时间戳相关性评分 - 不依赖固定格式"""
    import re
    
    score = 0
    
    # 1. 数字密度评分（更灵活）
    digit_count = sum(1 for c in text if c.isdigit())
    if digit_count > 0:
        digit_ratio = digit_count / len(text)
        if digit_ratio >= 0.8:  # 80%以上是数字
            score += 0.6
        elif digit_ratio >= 0.5:  # 50%以上是数字
            score += 0.4
        elif digit_ratio >= 0.3:  # 30%以上是数字
            score += 0.2
    
    # 2. 包含可能的时间分隔符（不限制特定组合）
    separators = [':', '-', '/', '.', ' ', '年', '月', '日', '时', '分', '秒']
    separator_count = sum(1 for sep in separators if sep in text)
    score += min(separator_count * 0.1, 0.3)
    
    # 3. 长度评分（监控时间戳的常见长度范围）
    text_len = len(text)
    if 10 <= text_len <= 25:  # 完整时间戳
        score += 0.3
    elif 6 <= text_len <= 9:   # 中等长度
        score += 0.25
    elif 3 <= text_len <= 5:   # 短数字（可能是部分时间戳）
        score += 0.2
    elif text_len >= 2:        # 最短的数字
        score += 0.15
    
    # 4. 数字连续性（时间戳通常包含连续数字）
    consecutive_digits = re.findall(r'\d+', text)
    if consecutive_digits:
        max_consecutive = max(len(num) for num in consecutive_digits)
        if max_consecutive >= 4:
            score += 0.2
        elif max_consecutive >= 3:
            score += 0.15
        elif max_consecutive >= 2:
            score += 0.1
    
    # 5. 常见时间数字范围检查（不强制格式）
    if text.isdigit() and len(text) >= 2:
        num = int(text)
        if 2000 <= num <= 2099:    # 年份范围
            score += 0.4
        elif 1 <= num <= 31:       # 日期范围
            score += 0.2
        elif 0 <= num <= 23:       # 小时范围
            score += 0.2
        elif 0 <= num <= 59:       # 分钟/秒钟范围  
            score += 0.2
        elif 100 <= num <= 999:    # 可能的年份片段或其他
            score += 0.15
    
    # 6. 包含中文时间词汇
    chinese_time_chars = ['年', '月', '日', '时', '分', '秒']
    if any(char in text for char in chinese_time_chars):
        score += 0.2
    
    return min(score, 1.0)  # 最大评分为1.0

def classify_timestamp_text(text):
    """分类时间戳文本类型"""
    import re
    
    if not text:
        return "empty"
    
    if text.isdigit():
        num = int(text)
        if 2000 <= num <= 2099:
            return "year"
        elif 1 <= num <= 31:
            return "day"
        elif 0 <= num <= 23:
            return "hour"
        elif 0 <= num <= 59:
            return "minute_second"
        else:
            return "number"
    
    if re.match(r'^\d{4}[-年/]\d{1,2}[-月/]\d{1,2}', text):
        return "date"
    elif re.match(r'^\d{1,2}[时:]\d{1,2}[分:]?\d{0,2}', text):
        return "time"
    elif any(sep in text for sep in ['年', '月', '日', '时', '分', '秒']):
        return "chinese_datetime"
    else:
        return "other"

def recognize_timestamp(image_path, ocr, use_preprocessing=False, min_confidence=0.2):
    """
    识别监控图像中的时间戳 - 优先使用原图，预处理作为备选
    """
    if not os.path.exists(image_path):
        return {"error": f"图片文件不存在: {image_path}"}
    
    try:
        print(f"正在识别时间戳: {os.path.basename(image_path)}")
        
        # 先用原始图像识别
        results = ocr.ocr(image_path, cls=True)
        
        # 如果原始图像识别不到结果且启用了预处理，则尝试预处理
        if use_preprocessing and (not results or not results[0]):
            original, processed = preprocess_timestamp_image(image_path)
            if processed is not None:
                # 保存预处理后的图像用于OCR
                temp_path = "temp_processed.jpg"
                cv2.imwrite(temp_path, processed)
                results = ocr.ocr(temp_path, cls=True)
                
                # 清理临时文件
                if os.path.exists(temp_path):
                    os.remove(temp_path)
        
        # 过滤和优化结果
        filtered_results = filter_timestamp_results(results, min_confidence)
        
        return {
            "success": True,
            "image_path": image_path,
            "total_detected": len(filtered_results),
            "results": filtered_results,
            "timestamp": datetime.now().isoformat()
        }
        
    except Exception as e:
        return {"error": f"时间戳识别失败: {str(e)}"}

def save_results_to_json(results, output_file="timestamp_results.json"):
    """保存识别结果到JSON文件"""
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"结果已保存到: {output_file}")
        return True
    except Exception as e:
        print(f"保存结果失败: {e}")
        return False

# 主程序 - 监控时间戳识别
def main():
    """主函数"""
    print("=" * 60)
    print("监控时间戳识别系统 - PaddleOCR版本")
    print("=" * 60)
    
    # 初始化PaddleOCR
    ocr, success = init_paddleocr()
    if not success:
        print("OCR引擎初始化失败，程序退出")
        return
    
    # 图片路径设置 - 支持命令行参数
    import sys
    if len(sys.argv) > 1:
        img_path = sys.argv[1]
    else:
        img_path = "C:\\Users\\KB\\Desktop\\Test.png"
    
    # 检查图片是否存在
    if not os.path.exists(img_path):
        print(f"图片文件不存在: {img_path}")
        print("请确认图片路径是否正确")
        return
    
    print(f"正在分析图片: {os.path.basename(img_path)}")
    
    # 识别时间戳
    result = recognize_timestamp(
        img_path, 
        ocr, 
        use_preprocessing=False,  # 默认不使用预处理，优先原图识别
        min_confidence=0.2  # 进一步降低置信度阈值，允许更多候选
    )
    
    # 处理结果
    if "error" in result:
        print(f"识别失败: {result['error']}")
        return
    
    # 显示结果
    print(f"\n检测到 {result['total_detected']} 个可能的时间戳文本")
    print("\n识别结果（按相关性排序）:")
    print("-" * 80)
    
    for i, item in enumerate(result['results'], 1):
        print(f"第{i}项:")
        print(f"  文字内容: {item['text']}")
        print(f"  OCR置信度: {item['confidence']:.3f}")
        print(f"  时间戳相关性: {item['timestamp_score']:.3f}")
        print(f"  综合得分: {item['total_score']:.3f}")
        if 'category' in item:
            print(f"  文本类型: {item['category']}")
        print(f"  坐标位置: {item['bbox']}")
        print("-" * 40)
    
    # 高亮显示最可能的时间戳
    if result['results']:
        best_match = result['results'][0]
        print("\n🎯 最可能的时间戳:")
        print(f"   📅 {best_match['text']}")
        print(f"   📊 综合得分: {best_match['total_score']:.3f}")
    
    # 保存结果
    output_file = f"timestamp_result_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    save_results_to_json(result, output_file)
    
    # 可视化结果（可选）
    try:
        visualize_results(img_path, result['results'])
    except Exception as e:
        print(f"可视化失败: {e}")
        print("跳过可视化步骤")
    
    print("\n✅ 时间戳识别完成！")

def visualize_results(image_path, results, show_top_n=3):
    """可视化识别结果"""
    try:
        # 读取图像
        image = Image.open(image_path).convert('RGB')
        
        # 创建图形
        plt.figure(figsize=(15, 10))
        plt.imshow(image)
        plt.axis('off')
        plt.title("监控时间戳识别结果", fontsize=16, pad=20)
        
        # 导入所需模块
        import matplotlib.patches as patches
        from matplotlib import rcParams
        
        # 设置中文字体（如果可用）
        try:
            rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
            rcParams['axes.unicode_minus'] = False
        except:
            pass
        
        ax = plt.gca()
        colors = ['red', 'blue', 'green', 'orange', 'purple']
        
        # 绘制检测框和文本
        for i, result in enumerate(results[:show_top_n]):
            bbox = result['bbox']
            text = result['text']
            score = result['total_score']
            
            # 转换坐标格式
            if len(bbox) == 4 and len(bbox[0]) == 2:
                # 四个角点格式
                bbox_array = np.array(bbox)
                polygon = patches.Polygon(
                    bbox_array, 
                    linewidth=3, 
                    edgecolor=colors[i % len(colors)], 
                    facecolor='none',
                    alpha=0.8
                )
                ax.add_patch(polygon)
                
                # 添加文本标签
                x_center = np.mean(bbox_array[:, 0])
                y_center = np.mean(bbox_array[:, 1]) - 15
                
            plt.text(
                x_center, y_center, 
                f'#{i+1}: {text}\n得分: {score:.2f}', 
                fontsize=10, 
                color=colors[i % len(colors)],
                ha='center',
                va='bottom',
                bbox=dict(
                    boxstyle="round,pad=0.5", 
                    facecolor="white", 
                    alpha=0.9,
                    edgecolor=colors[i % len(colors)]
                )
            )
        
        plt.tight_layout()
        plt.show()
        
    except Exception as e:
        print(f"可视化过程出错: {e}")

if __name__ == "__main__":
    main()
